from PIL import Image, UnidentifiedImageError
import io
from typing import Dict, Optional, Tuple
import logging
import base64
import struct

class ImageProcessor:
    """图片处理器，用于处理和压缩图片
    
    支持的功能：
    - 图片格式转换
    - 图片压缩
    - 尺寸调整
    - 保持透明通道（可选）
    """
    
    def __init__(self, config: dict):
        self.config = config
        self.max_size = config['document']['image']['max_size']
        self.quality = config['document']['image']['quality']
        self.supported_formats = config['document']['image']['supported_formats']
        
    def _compress_image(self, img: Image.Image, max_size: int) -> Tuple[Image.Image, bytes]:
        """压缩图片到指定大小
        
        Args:
            img: 原始图片
            max_size: 最大允许大小（字节）
            
        Returns:
            压缩后的图片对象和二进制数据
        """
        output = io.BytesIO()
        img.save(output, format='JPEG', quality=self.quality, optimize=True)
        
        while output.tell() > max_size and img.size[0] > 800:
            new_size = (int(img.size[0] * 0.8), int(img.size[1] * 0.8))
            img = img.resize(new_size, Image.Resampling.LANCZOS)
            output = io.BytesIO()
            img.save(output, format='JPEG', quality=self.quality, optimize=True)
            logging.info(f"压缩后大小: {output.tell()}字节")
            
        return img, output.getvalue()

    def process(self, image_data: Dict) -> Optional[Dict]:
        """处理图片数据
        
        Args:
            image_data: 包含图片数据的字典，必须包含'image_data'键
            
        Returns:
            处理后的图片数据字典，失败返回None
        """
        try:
            # 获取二进制数据和格式
            binary_data = image_data.get('image_data') if isinstance(image_data, dict) else image_data
            format = image_data.get('format', 'JPEG').upper() if isinstance(image_data, dict) else 'JPEG'
            
            if not binary_data:
                logging.error("图片数据为空")
                return None
                
            logging.info(f"处理图片: 原始格式={format}, 大小={len(binary_data)}字节")
            
            # 使用BytesIO包装二进制数据
            input_stream = io.BytesIO(binary_data)
            
            try:
                # 打开并验证图片
                img = Image.open(input_stream)
                img.verify()
                input_stream.seek(0)
                img = Image.open(input_stream)
                
                # 处理透明通道
                if img.mode in ('RGBA', 'LA') or (img.mode == 'P' and 'transparency' in img.info):
                    # 创建白色背景
                    background = Image.new('RGB', img.size, (255, 255, 255))
                    if img.mode == 'P':
                        img = img.convert('RGBA')
                    # 合并图层
                    background.paste(img, mask=img.split()[3] if img.mode == 'RGBA' else img.split()[1])
                    img = background
                elif img.mode != 'RGB':
                    img = img.convert('RGB')
                
                # 压缩图片
                img, processed_data = self._compress_image(img, self.max_size)
                
                # 更新图片数据
                if isinstance(image_data, dict):
                    image_data['image_data'] = processed_data
                    image_data['format'] = 'jpeg'
                else:
                    image_data = {
                        'image_data': processed_data,
                        'format': 'jpeg'
                    }
                
                logging.info(f"处理完成: 大小={len(image_data['image_data'])}字节")
                return image_data
                
            except UnidentifiedImageError:
                logging.error(f"无法识别图片格式: {format}")
                return None
                
            except Exception as e:
                logging.error(f"处理图片失败: {str(e)}")
                return None
                
        except Exception as e:
            logging.error(f"处理图片时发生异常: {str(e)}")
            return None

    def _get_save_format(self, format_str):
        """获取保存格式（已弃用，统一使用JPEG）"""
        return 'JPEG' 